# -*- coding: utf-8 -*-
"""
从第1行到第49行这部分python代码，对应的是matlab的第1行到第49行
对应Matlab文件: Func_Extract_Field.m
功能: 提取场强曲线并换算为接收功率(dBm)
"""

import numpy as np


def func_extract_field(f, distance, TX, RX, shape, cross_section, dx, ds, E):
    """
    从第1行到第49行这部分python代码，对应的是matlab的第1行到第49行
    Matlab: function [P] = Func_Extract_Field(f, distance, TX, RX, shape, cross_section, dx, ds, E)
    返回: P (沿轴向采样的接收功率dBm)
    """
    # 从第2行到第18行，对应matlab的第2行到第18行 (形状参数解包)
    if shape == 1:
        H_top = cross_section[1] / 2
        W_half = cross_section[0] / 2
        H_bottom = H_top
    elif shape == 2:
        H_top = cross_section[1]
        W_half = cross_section[0]
        H_bottom = cross_section[3]
    elif shape == 3:
        H_top = cross_section[1]
        W_half = cross_section[0]
        H_bottom = cross_section[3]
    elif shape == 4:
        H_top = cross_section[1]
        W_half = cross_section[0]
        H_bottom = cross_section[3]
    else:
        raise ValueError("Unsupported shape value: {}".format(shape))

    # 从第20行到第28行，对应matlab的第20行到第28行 (网格参数)
    dy = dx
    nh_top = int(np.ceil(H_top / dy))
    nw_half = int(np.ceil(W_half / dx))
    Nx = 2 * nw_half
    Nz = int(np.ceil(distance / ds))
    x_c = Nx / 2 + 1
    y_c = nh_top + 1

    # 从第30行到第36行，对应matlab的第30行到第36行 (接收点位置)
    Rx_horizontal = RX[0]
    Rx_height = RX[1]
    Rx_vertical = Rx_height - H_bottom
    rp_y = int(y_c - np.round(Rx_vertical / dy))
    rp_x = int(x_c + np.round(Rx_horizontal / dx))
    E_extract = np.zeros((Nz + 1,), dtype=complex)
    # 注意MATLAB索引是(y, x, z)，Python/numpy同样采用E[rp_y, rp_x, :]
    E_extract[:] = E[rp_y, rp_x, :]

    # 从第38行到第49行，对应matlab的第38行到第49行 (功率换算)
    c = 299792458.0
    lambda_ = c / f
    mu_0 = 4 * np.pi * 1e-7
    eps_0 = 1 / mu_0 / c / c
    eita = 1 / np.sqrt(eps_0 / mu_0)
    l_eff = lambda_ / np.pi
    Rs = 73
    E_dB = 10 * np.log10((l_eff * np.abs(E_extract)) ** 2 / 8 / Rs)
    Tx_Gain_Diff = (TX[3] - 30) + 10 * np.log(eita / 2 / np.pi) - TX[2]
    Rx_Gain_Diff = RX[2] - 2.15
    VPE_shift = 30
    P = E_dB + Tx_Gain_Diff + Rx_Gain_Diff + VPE_shift
    # 对应MATLAB: P = P(1:round(RX{4}/ds):end);
    step = int(np.round(RX[3] / ds))
    P = P[::step]
    return P
